package club.ruanx.spider.filter;

import club.ruanx.spider.entity.SpiderRecord;
import club.ruanx.spider.repository.SpiderRecordRepository;
import com.google.common.hash.BloomFilter;
import com.google.common.hash.Funnels;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.nio.charset.Charset;
import java.util.List;

/**
 * @author 阮胜
 * @date 2018/11/26 13:24
 */
@SuppressWarnings("UnstableApiUsage")
@Component
@Slf4j
public class SpiderUrlFilter {
    private final BloomFilter<String> bloomFilter;

    public SpiderUrlFilter(SpiderRecordRepository spiderRecordRepository) {
        List<SpiderRecord> spiderRecordList = spiderRecordRepository.findAll();
        bloomFilter = BloomFilter.create(Funnels.stringFunnel(Charset.defaultCharset()), spiderRecordList.size() * 5);
        spiderRecordList.forEach(spiderRecord -> bloomFilter.put(spiderRecord.getUrl()));
        log.info("初始化布隆过滤器成功,共初始化{}条URL.", spiderRecordList.size());
    }

    /**
     * 向过滤器添加要过滤的URL,如果已经存在该URL False
     *
     * @param url
     * @return
     */
    public boolean add(String url) {
        return bloomFilter.put(url);
    }

    public void add(Iterable<String> urls) {
        urls.forEach(bloomFilter::put);
    }

    public boolean contains(String url) {
        return bloomFilter.mightContain(url);
    }
}
